import torch
import numpy as np
from sklearn.datasets import load_boston
from sklearn.preprocessing import MinMaxScaler
from sklearn.model_selection import train_test_split

# 加载波士顿房价数据集
boston = load_boston()
data = boston.data
target = boston.target.reshape(-1, 1)

# 数据标准化
scaler_data = MinMaxScaler()
scaler_target = MinMaxScaler()
data = scaler_data.fit_transform(data)
target = scaler_target.fit_transform(target)

# 定义时间窗口大小
c = 7
x = []
y = []
for i in range(len(data) - c):
    x_data = data[i:i + c]
    y_data = target[i + c]
    x.append(x_data)
    y.append(y_data)

x = torch.Tensor(x)
y = torch.Tensor(y)

# 划分训练集和测试集
train_x, test_x, train_y, test_y = train_test_split(x, y, shuffle=False)
print(train_x.shape)

class RNN(torch.nn.Module):
    def __init__(self) -> None:
        super().__init__()
        self.rnn = torch.nn.RNN(input_size=13, hidden_size=50, batch_first=True)
        self.fc = torch.nn.Linear(in_features=50, out_features=1)

    def forward(self, x):
        x, _ = self.rnn(x)
        out = self.fc(x[:, -1, :])
        return out

model = RNN()
loss_fn = torch.nn.MSELoss()
op = torch.optim.Adam(model.parameters(), lr=0.01)
loss_list = []
for i in range(1000):
    op.zero_grad()
    h = model(train_x)
    loss = loss_fn(h, train_y)
    loss.backward()
    loss_list.append(loss.item())
    op.step()

# 进行预测
pre = model(test_x).reshape(-1)
print(pre)

# 可视化预测结果和真实值
import matplotlib.pyplot as plt

plt.plot(pre.data.numpy(), c='r', label='Prediction')
plt.plot(test_y.data.numpy(), c='g', label='Actual')
plt.legend()
plt.show()